from typing import TypedDict

from langgraph.graph import StateGraph

# 1. 首先，定义一个状态结构（State）
class MyState(TypedDict):
    input: str
    result: str

# 2. 创建一个基于该状态的工作流构建器
builder = StateGraph(MyState)

# 3. 定义多个节点（Node），并将它们添加到构建器中
def node1_function(state: MyState):
    print("node1 executed")
    return {"result": f"Processed: {state['input']}"}

builder.add_node("node_1", node1_function)

def node2_function(state: MyState):
    print("node2 executed")
    return {"result": state['result'].upper()}

builder.add_node("node_2", node2_function)

# 4. 【关键步骤】设置入口点
builder.set_entry_point("node_1") # 告诉图，从 "node_1" 开始执行

# 5. 定义节点之间的连接关系（边）
builder.add_edge("node_1", "node_2") # 执行完 node_1 后，前往 node_2

# 6. 设置终点（可选，但推荐）
builder.set_finish_point("node_2") # 告诉图，node_2 是终点之一

# 7. 编译图，使其可执行
graph = builder.compile()

graph.invoke({"input": "start workflow"})